
#include "MatrixSolver.hpp"
#include "Util.hpp"
#include "VAR.hpp"
#include <stdexcept>
#include <fstream>
#include <cstdlib>
#include <unistd.h>

using namespace std;
using namespace Util;

int
main ( int argc, char* argv[], char* env[] )
{
  // defaults
  bool HasConst = false;
  bool test = false;
  int nStep = 5;
  int nData = 0;
  int seed = 0;
  int c;
  while ((c = getopt(argc, argv, "cn:s:t")) != EOF)
    switch (c)
      { 
      case 'c':
	HasConst = true;
        break;
      case 'n':
        nStep = atoi ( optarg );
        break;
      case 's':
        seed = atoi ( optarg );
        break;
      case 't':
	test = true;
        break;
      case '?':
        cerr << "invalid option" << endl;
        exit(1);
      }

  //
  // Initialize RNG
  //
  if ( seed > 0 )
    seed = -seed;
  Normal* norm = new Normal(seed);

  if ( argc - optind < 1 )
    {
      cerr << "Usage: " << argv[0] << " <VAR input>" << endl;
      return 1;
    }

  try
    {
      int order, dim;

      // Z transpose - complete data vector
      Matrix<double> Zt;
      // Yt - accumulation
      Matrix<double> Yt_acc;

      ofstream VAR_out("VAR.out");
      if ( VAR_out.fail() )
	{
	  cerr << "open \"sim.txt\" failed" << endl;
	  exit(1);
	}

      if ( test )
	{
	  //
	  // input data
	  //
	  ifstream VAR_in(argv[optind]);
	  if ( VAR_in.fail() )
	    {
	      cerr << "open " << argv[optind] << " failed" << endl;
	      exit(1);
	    }
	  VAR_in >> order >> dim;

	  cout << "VAR for order " << order << ", dimension " << dim << endl;
	  VAR* iVAR = new VAR(order,dim);
	  //
	  // read coefficient matrices
	  //
	  Matrix<double> At(dim,dim);
	  for ( int i=0; i < order; ++i )
	    {
	      VAR_in >> At;
	      iVAR->Ai(At,i);
	    }
	  // const coeff
	  Vector<double> Cvec(dim);
	  VAR_in >> Cvec;
	  iVAR->C(Cvec);
	  //
	  // read initial state vectors
	  //
	  Vector<double> Yt(dim);
	  for ( int i=0; i < order; ++i )
	    {
	      VAR_in >> Yt;
	      iVAR->Yi(Yt,i);
	    }
	  //
	  // impulse
	  //
	  Vector<double> St(dim);
	  VAR_in >> St;
	  iVAR->Sigma(St);
	  // end input
	  VAR_in.close();

	  //
	  // echo input data
	  //
	  VAR_out << "VAR for order " << order << ", dimension " << dim << endl;
	  for ( int i=0; i < order; ++i )
	    VAR_out << iVAR->Ai(i);

	  for ( int i=0; i < order; ++i )
	    VAR_out << iVAR->Yi(i);
	  VAR_out << St;

	  //
	  // simulate
	  //

	  int Ys_size = dim * order;
	  int i0 = 0;
	  // lengthen for a a constant
	  if ( HasConst )
	    {
	      i0 = 1;
	      Ys_size += i0;
	    }
	  Vector<double> Ys ( Ys_size );
	  // this should never be overwitten
	  Ys[0] = 1.0;

	  //
	  // Simulate to produce test data
	  //
	  // disturbance matrix
	  //Matrix<double> Ed;
	  //Vector<double> Et;
	  VAR_out << "simulation" << endl;
	  for ( int k=0; k < nStep; ++k )
	    {
	      // collect data
	      for ( int j=0; j < order; ++j )
		{
		  Yt = iVAR->Yi(j);
		  for ( int i=0; i < dim; ++i )
		    Ys[j*dim+i+i0] = Yt[i];
		}
	      Zt.push_back(Ys);
	      Yt = iVAR->simulate();
	      //Et = iVAR->E();
	      Yt_acc.push_back(Yt);
	      //Ed.push_back(Et);
	      //VAR_out << Et;
	      VAR_out << Yt;
	    }      
	  //
	  // end test section
	  //
	} 
      else 
	{
	  //
	  // read some real data
	  //
	  //
	  // input data
	  //
	  ifstream VAR_in(argv[optind]);
	  if ( VAR_in.fail() )
	    {
	      cerr << "open " << argv[optind] << " failed" << endl;
	      exit(1);
	    }
	  cout << "what's the desired order? > ";
	  cin >> order;
	  cout << endl;
	  // how many columns of data?
	  string stmp,line;
	  Vector<string> Header;
	  getline (VAR_in,line);
	  istringstream istr(line);
	  while ( istr >> stmp )
	    Header.push_back(stmp);
	  dim = Header.size();
	  cout << "dimension: " << dim << endl;
	  Vector<double> Yt(dim);

	  Matrix<double> Yt_total;

	  while ( ! VAR_in.eof() )
	    {
	      getline (VAR_in,line);
	      istringstream iss(line);
	      int i = 0;
	      while ( iss >> Yt[i] )
		++i;
	      if ( ! i )
		break;
	      Yt_total.push_back(Yt);
	    }
	  nData = Yt_total.row();
	  cout << nData << " rows of data" << endl;

	  int Ys_size = dim * order;
	  int i0 = 0;
	  // lengthen for a constant
	  if ( HasConst )
	    {
	      i0 = 1;
	      Ys_size += i0;
	    }
	  Vector<double> Ys ( Ys_size );
	  // this should never be overwitten
	  Ys[0] = 1.0;

	  // Fill Zt
	  for ( int j=order; j < nData-1; ++j )
	    {
	      for ( int k=0; k < order; ++k )
		{
		  Yt = Yt_total[j-k];
		  for ( int i=0; i < dim; ++i )
		    Ys[k*dim+i+i0] = Yt[i];
		}
	      Zt.push_back(Ys);
	      Yt_acc.push_back(Yt_total[j+1]);
	    }
	}

      //
      // Linear Regression
      //
      // transpose of the transpose
      Matrix<double> Z = Zt.Tr();
      // full-data matrix
      MatrixSolver<double> ZZt = Z * Zt;
      //cout << "ZZt: " << ZZt.row() << ", " << ZZt.col() << endl;
      
      //cout << "ZZt\n" << ZZt;
      // ZZt is destroyed
      Matrix<double>Cinv = ZZt.Cholesky_Inverse();
      //cout << "ZZt Inv\n" << Cinv;

      Matrix<double> Mtmp = Zt * Cinv;
      //cout << "Yt_acc: " << Yt_acc.row() << ", " << Yt_acc.col() << endl;
      //cout << "Mtmp: " << Mtmp.row() << ", " << Mtmp.col() << endl;
      Matrix<double> Y = Yt_acc.Tr();
      Matrix<double> B_hat =  Y * Mtmp;
      VAR_out << "Coefficient Estimate" << endl;
      B_hat.IOS_Format(ios::scientific);
      VAR_out << B_hat;

      Matrix<double> Y_hat = B_hat * Z;

      Vector<double> V1(Y.col(),1.0);
      Vector<double> Y_mean = Y * V1;
      Y_mean /= V1.size();

      Matrix<double> Y_diff = Y - Y_mean;
      MatrixSolver<double> T_SS = Y_diff * Y_diff.Tr();
      VAR_out << "Total Sum Squared" << endl;
      VAR_out << T_SS;
      
      Y_diff = Y_hat - Y_mean;
      Matrix<double> E_SS = Y_diff * Y_diff.Tr();
      VAR_out << "Explained Sum Squared" << endl;
      VAR_out << E_SS;

      // Residual
      Matrix<double>Y_res = Y - Y_hat;
      //cout << "Residual Matrix: " << Y_res.row() << ", " << Y_res.col() << endl;
      Matrix<double> R_SS = Y_res * Y_res.Tr();
      VAR_out << "Residual Sum Squared" << endl;
      R_SS.IOS_Format(ios::scientific);
      VAR_out << R_SS;

      Matrix<double> StdErr = R_SS / ( Y_res.col() - 2 );
      VAR_out << "Standard Error" << endl;
      StdErr.IOS_Format(ios::scientific);
      VAR_out << StdErr;

      Vector<double> Z_mean = Z * V1;
      //if ( HasConst )
      //Z_mean[0] = 0;
      Z_mean /= V1.size();
      Matrix<double> Z_diff = Z - Z_mean;
      // MatrixSolver<double> Z_SS = Z_diff * Z_diff.Tr();
      // cout << "Z_SS: " << Z_SS.row() << ", " << Z_SS.col() << endl;
      // cout << Z_SS;

      //
      // R^2
      //
      Matrix<double> BZ = B_hat * Z_diff;
      Matrix<double> BZZBt = BZ * BZ.Tr();
      Matrix<double> T_SS_Inv = T_SS.Cholesky_Inverse();
      Matrix<double> R2 = T_SS_Inv * BZZBt;
      VAR_out << "R^2" << endl;
      R2.IOS_Format(ios::scientific);
      VAR_out << R2;

      //Matrix<double> Z_SS_Inv = Z_SS.Cholesky_Inverse();
      //
      // StdErr for the Coeffs
      //
      // Matrix<double> B_se = B_hat * Zss_Inv;
      //cout << "Z_SS Inverse" << endl;
      //cout << Z_SS_Inv;

      //
      // End Linear Regression
      //

      // covariance of disturbance
      //Matrix<double> M_S = Ed.Tr() * Ed;
      //M_S /= Ed.row();
      //cout << "Dist Covar: " << M_S.row() << ", " << M_S.col() << endl;

      //VAR_out << "Parameter Estimate" << endl;
      //VAR_out << B_hat;

      //VAR_out << "Disturbance Covar" << endl;
      //VAR_out << M_S;
      VAR_out.close();
    }
  catch ( std::exception& ex )
    {
      cerr << ex.what() << endl;
      exit(1);
    }
  exit (0);
}
